# -*- coding: utf-8 -*-

from xml.dom import minidom
import os, sys, glob, shutil
from tqdm import tqdm as tqdm

file_path = os.path.abspath(__file__)
sys.path.append(os.path.abspath(os.path.join(file_path, "..", "..", "..")))
from code_aculat.data_convert.voc2coco import get_categories
from code_aculat.data_analyse.data_analyse_coco import analyse_num_each_class_drop_rare_cat


def convert_coordinates(size, box):
    dw = 1.0 / size[0]
    dh = 1.0 / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def check_txt():
    import shutil
    txt_dir = r"/Users/edy/Downloads/get_drink/get_drink_from_fd/return/txt"
    img_dir = r"/Users/edy/Downloads/get_drink/get_drink_from_fd/return/img"

    for txt in filter(lambda x: x.endswith('.txt'), os.listdir(txt_dir)):
        with open(os.path.join(txt_dir, txt), 'r') as f:
            record = f.readlines()
        if len(record) < 1:
            os.remove(os.path.join(txt_dir, txt))
            os.remove(os.path.join(img_dir, txt.replace('.txt', '.jpg')))


def move_rare_category_txt():
    "移动包含指定类别的数据,目前看不太好弄，一移动验证集里都没有数据了"
    json_path = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/voc2coco.json"
    txt_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/yolo_txt_3"
    catename_to_id_txt_path = os.path.join(txt_dir, 'class.txt')
    origin_txt_dir = os.path.join(txt_dir, 'test')
    dst_txt_dir = os.path.join(txt_dir, 'train')

    operate_txt = []
    all_catename_to_id = {}
    # 获取到这些类别
    _, catename = analyse_num_each_class_drop_rare_cat(json_path, False)
    # 把类名转成catid
    with open(catename_to_id_txt_path, 'r') as f:
        records = f.readlines()

    for rec in records:
        name, id = rec.strip('\n').split("  ")  # 两个英文空格
        all_catename_to_id[name] = id

    drop_cate_id = [all_catename_to_id[x] for x in catename]

    # 遍历所有val里的txt文件,找出所有包含这些类别的txt，然后全移动到train里
    for txt in glob.glob(os.path.join(origin_txt_dir, "*.txt")):
        with open(txt, 'r') as f:
            txt_obs = f.readlines()

        for obs in txt_obs:
            cate_id, _, _, _, _ = obs.strip('\n').split(' ')
            if cate_id in drop_cate_id:
                operate_txt.append(os.path.basename(txt))
                break

    # 移动
    for txt_m in operate_txt:
        shutil.move(os.path.join(origin_txt_dir, txt_m), os.path.join(dst_txt_dir, txt_m))


def move_txt_for_category_absence(txt_dir, cat_num):
    "训练集里少了某个类的目标，从验证集移出去, txt_dir到labels目录"
    origin_txt_dir = os.path.join(txt_dir, 'val')
    dst_txt_dir = os.path.join(txt_dir, 'train')

    # cat_num = 348
    all_cate_id = [x for x in range(cat_num)]
    target_txt = []
    move_num = 0

    # 遍历训练集的txt文件，找到缺失的类
    train_cate_id = []
    for txt in glob.glob(os.path.join(dst_txt_dir, "*.txt")):
        with open(txt, 'r') as f:
            txt_obs = f.readlines()

        for obs in txt_obs:
            cate_id, _, _, _, _ = obs.strip('\n').split(' ')
            cate_id = int(cate_id)
            if cate_id not in train_cate_id:
                train_cate_id.append(cate_id)

    absence_catid = [x for x in all_cate_id if x not in train_cate_id]
    # 遍历验证集的txt文件，找到缺失类的txt
    if len(absence_catid):
        for txt in glob.glob(os.path.join(origin_txt_dir, "*.txt")):
            with open(txt, 'r') as f:
                txt_obs = f.readlines()

            for obs in txt_obs:
                cate_id, _, _, _, _ = obs.strip('\n').split(' ')
                cate_id = int(cate_id)
                if cate_id in absence_catid and os.path.basename(txt) not in target_txt:
                    target_txt.append(os.path.basename(txt))
    # 把验证集的txt移动到训练集
    for txt_m in target_txt:
        shutil.move(os.path.join(origin_txt_dir, txt_m), os.path.join(dst_txt_dir, txt_m))
        move_num += 1

    print("move txt file number is %d" % move_num)


def convert_xml2yolo(xml_files, out_txt_dir, PRE_DEFINE_CATEGORIES):
    # lut = {}
    # lut["position"] = 1  # 类别及其cate_id
    # xml_dir = r"/Users/edy/Downloads/get_drink/get_drink_from_fd/return/xml"
    # out_txt_dir = r"/Users/edy/Downloads/get_drink/get_drink_from_fd/return/txt"
    # os.makedirs(out_txt_dir)

    for fname in tqdm(xml_files):

        xmldoc = minidom.parse(fname)
        txt_name = os.path.basename(fname).split('.')[0] + '.txt'
        fname_out = os.path.join(out_txt_dir, txt_name)

        with open(fname_out, "w") as f:

            itemlist = xmldoc.getElementsByTagName('object')
            size = xmldoc.getElementsByTagName('size')[0]
            width = int((size.getElementsByTagName('width')[0]).firstChild.data)
            height = int((size.getElementsByTagName('height')[0]).firstChild.data)

            for item in itemlist:
                # get class label
                classid = (item.getElementsByTagName('name')[0]).firstChild.data
                # if classid in ["row",'bottle']:  # 去掉围栏
                #     continue
                if classid in PRE_DEFINE_CATEGORIES:
                    label_str = str(PRE_DEFINE_CATEGORIES[classid])
                else:
                    label_str = str(PRE_DEFINE_CATEGORIES["bottle"]) # 不在范围内的类视为bottle吧
                    # continue  # 不在范围内的类过滤掉


                # get bbox coordinates
                xmin = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('xmin')[0]).firstChild.data
                ymin = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('ymin')[0]).firstChild.data
                xmax = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('xmax')[0]).firstChild.data
                ymax = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('ymax')[0]).firstChild.data

                ##for track
                # b = (float(xmin), float(ymin), float(xmax), float(ymax))
                # f.write(label_str + " " + " ".join([("%.6f" % a) for a in b]) + '\n')

                b = (float(xmin), float(xmax), float(ymin), float(ymax))  # for yolo
                bb = convert_coordinates((width, height), b)

                f.write(label_str + " " + " ".join([("%.6f" % a) for a in bb]) + '\n')

        # print("wrote %s" % fname_out)


def convert(xml_dir, csv_path, txt_dir, PRE_DEFINE_CATEGORIES):
    if csv_path:
        with open(csv_path, "r") as f:
            lines_record = f.readlines()
            lines_record.pop(0)

        xml_files = []

        for line in lines_record:
            xml_files.append(os.path.join(xml_dir, line.strip("\n").split(',')[1] + ".xml"))

    else:
        xml_files = [os.path.join(xml_dir, file) for file in os.listdir(xml_dir)]

    convert_xml2yolo(xml_files, txt_dir, PRE_DEFINE_CATEGORIES)


def copy_images_from_txt():
    origin_image_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/images/train"
    out_image_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/out_split/folds/fold_v1/images"
    img_format = ".jpg"

    txt_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/out_split/folds/fold_v1/labels"

    for dir in os.listdir(txt_dir):
        if not os.path.isdir(os.path.join(txt_dir, dir)):
            continue

        dst_img_dir = os.path.join(out_image_dir, dir)
        os.makedirs(dst_img_dir)

        for txt in tqdm(filter(lambda x: x.endswith('.txt'), os.listdir(os.path.join(txt_dir, dir)))):
            img_name = txt.replace(".txt", img_format)
            shutil.copy(os.path.join(origin_image_dir, img_name), os.path.join(dst_img_dir, img_name))


def main():
    csv_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/out_split/folds/fold_v1"
    xml_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/Annotations"
    txt_out_dir = r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/out_split/folds/fold_v1/labels"

    os.makedirs(txt_out_dir, exist_ok=True)
    name_to_id_txt_path = os.path.join(txt_out_dir, 'class.txt')
    PRE_DEFINE_CATEGORIES = None
    if PRE_DEFINE_CATEGORIES is None:
        PRE_DEFINE_CATEGORIES = get_categories(glob.glob(os.path.join(xml_dir, '*.xml')))

    # 存储类名和caid的对应字典
    with open(name_to_id_txt_path, 'w', encoding='utf-8') as f:
        for name, id in PRE_DEFINE_CATEGORIES.items():
            f.write("%s  %d\n" % (name, id))

    # If you want to do train/test split, you can pass a subset of xml files to convert function.
    if csv_dir == None:
        csv_files = None

    elif os.path.isdir(csv_dir):
        csv_files = os.listdir(csv_dir)
        csv_files = list(filter(lambda x: x.endswith(".csv"), csv_files))
        csv_files = [os.path.join(csv_dir, csv) for csv in csv_files]

    elif os.path.isfile(csv_dir):
        csv_files = [csv_dir]

    else:
        raise ValueError

    if csv_files:

        for csv_path in csv_files:
            txt_dir = os.path.join(txt_out_dir, os.path.basename(csv_path).split('.')[0])
            os.makedirs(txt_dir)

            # 从csv里获取xml并进行转换
            convert(xml_dir, csv_path, txt_dir, PRE_DEFINE_CATEGORIES)
    else:
        os.makedirs(txt_out_dir)
        # 没有csv文件，则直接将所有的xml转成txt
        convert(xml_dir, csv_files, txt_out_dir, PRE_DEFINE_CATEGORIES)
    print("Success: {}".format(txt_out_dir))


if __name__ == '__main__':
    # main()
    # move_rare_category_txt()
    copy_images_from_txt()
    # move_txt_for_category_absence()
